#!/bin/env python

import cv2
import numpy as np
import pickle
import gzip

#读取mnist.pkl.gz数据
def load_data():
    mnist = gzip.open( '../../resource/data/MNIST/mnist.pkl.gz', 'rb')
    #训练集，校验集和测试集 注意 pickle.load(mnist,encoding="bytes")
    training_data, classification_data, test_data = pickle.load(mnist, encoding= "bytes")
    mnist.close()
    return (training_data, classification_data, test_data)

def wrap_data():
    tr_d, va_d, te_d = load_data()

    training_inputs = tr_d[0]
    training_results = [vectorized_result(y) for y in tr_d[1]]
    training_data = zip(training_inputs, training_results)

    validation_inputs = va_d[0]
    validation_data = zip(va_d[0], va_d[1])

    test_inputs = te_d[0]
    test_data = zip(te_d[0], te_d[1])

    print("训练集样本数:", len(training_inputs))
    print("检验集样本数:", len(validation_inputs))
    print("测试集样本数:", len(test_inputs))

    #训练集，校验集和测试集，只有训练集是有results的，也就是由监督的
    return (training_data, validation_data, test_data)

#创建包含10个元素的零元组，在期望结果的位置设置1.这样可以用作输出层的类标签
def vectorized_result(j):
    e = np.zeros((10,), np.float32)
    e[j] = 1.0
    return e

# 创建一个手写数字识别的ANN
def create_ANN(hidden = 20):
    # 创建人工神经网络
    ann = cv2.ml.ANN_MLP_create()
    # 创建拓扑结构
    ann.setLayerSizes(np.array([784, hidden, 10]))
    # 设置训练方法
    ann.setTrainMethod(cv2.ml.ANN_MLP_BACKPROP, 0.1, 0.1)
    # 设置激活函数
    ann.setActivationFunction(cv2.ml.ANN_MLP_SIGMOID_SYM, 0.6, 1.0)
    # 设置终止条件
    ann.setTermCriteria((cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 100, 1.0))
    return ann

# 训练函数
def train(ann, samples = 10000, epochs = 1):
    # 给定一定数量的样本和训练周期，加载数据，然后迭代某个设定的次数
    tr, val, test = wrap_data()
    tr = list(tr)
    for x in range(epochs):
        counter = 0
        for img in tr:
            if (counter > samples):
                break
            if (counter % 1000 == 0):
                print( "Epoch %d : Trained %d / %d " % (x, counter, samples))
            counter += 1
            data, digit = img
            train_data_set = cv2.ml.TrainData_create(np.array([data], dtype=np.float32), cv2.ml.ROW_SAMPLE, np.array([digit], dtype=np.float32))
            if ann.isTrained():
                ann.train(train_data_set, cv2.ml.ANN_MLP_UPDATE_WEIGHTS | cv2.ml.ANN_MLP_NO_INPUT_SCALE | cv2.ml.ANN_MLP_NO_OUTPUT_SCALE)
            else:
                ann.train(train_data_set, cv2.ml.ANN_MLP_NO_INPUT_SCALE | cv2.ml.ANN_MLP_NO_OUTPUT_SCALE)
        print( "Epoch %d complete" % x)
    return ann, test

#封装ann的test
def test(ann, test_data):
    sample = np.array(test_data[0], dtype=np.float32)
    label = test_data[1]
    image = sample.reshape(28, 28)
    predict = ann.predict(np.array([sample], np.float32))

    return (image, label, int(predict[0]))

#封装ann的predict
def predict(ann, sample):
    resized = sample.copy()
    rows, cols = resized.shape
    if (rows != 28 or cols != 28) and rows * cols > 0:
        resized = cv2.resize(resized, (28, 28), interpolation = cv2.INTER_LINEAR)
    return ann.predict(np.array([resized.ravel()], dtype=np.float32))



